import numpy as np
import matplotlib.pyplot as plt
from scipy.special import comb

"""
a, b: params of cost function
zc: z-axis intersection
Tsup: support time
sx, sy: walk parameters(vectors)
xi, yi: initial CoM coordinate
px, py: initial CoM foot placement
vxi, vyi: initial CoM velocity
theta: changing degrees (not rad!) per step
"""
n = 0  # step number
g = 9.81  # gravity: m/s^2

a = 1
b = 1
zc = 0.3
Tsup = 2.0

sx = [0.0, 0.15, 0.15, 0.15, 0, 0]
sy = [0.05, 0.08, 0.05, 0.05, 0.05, 0]
theta = [0, 20, 40, 60, 60, 60]
px = 0
py = 0
xi = 0
yi = 0
vxi = 0
vyi = 0
xt = 0
yt = 0
vxt = 0
vyt = 0

"""The step after the last step should be zero"""
# sx.append(0)
# sy.append(0)

"""Desired foot placement"""
px0 = px
py0 = py

"""Calculate basic paras"""
Tc = np.sqrt(zc / g)
C = np.cosh(Tsup / Tc)
S = np.sinh(Tsup / Tc)
D = a * (C - 1) ** 2 + b * (S / Tc) ** 2
radian = np.deg2rad(theta)
st = np.sin(radian)
ct = np.cos(radian)


def get_value(name):
    return eval(name)


def rotate(i):
    if i == 0:
        R = np.eye(2)
    else:
        R = np.array([[ct[i], -st[i]], [st[i], ct[i]]])
    return R


def cal_foot_place():
    global px0, py0, n
    if n != 0:
        pxy0 = [[px0], [py0]]
        sxy = [[sx[n]], [-(-1) ** n * sy[n]]]
        pxy0 = pxy0 + rotate(n).dot(sxy)

        px0 = pxy0[0][0]
        py0 = pxy0[1][0]
        # print(px0)
        # print(py0)

    """Calculate the coordinate (xbar, ybar)"""
    xybar = [[sx[n + 1] / 2], [-(-1) ** n * sy[n] / 2]]
    xybar = rotate(n + 1).dot(xybar)
    xbar = xybar[0][0]
    ybar = xybar[1][0]
    vxybar = [[(C + 1) / (Tc * S) * xbar], [(C - 1) / (Tc * S) * ybar]]
    vxybar = rotate(0).dot(vxybar)
    vxbar = vxybar[0][0]
    vybar = vxybar[1][0]

    """Target state of CoM of hte n-th step"""
    xd = px0 + xbar
    yd = py0 + ybar
    vxd = vxbar
    vyd = vybar
    # print(xd)
    # print(vxd)

    """
    Update px, py to be the real foot place in step n,
    a, b are parameters for cost function
    """
    global px, py
    px = -a * (C - 1) / D * (xd - C * xi - Tc * S * vxi) - b * S / (Tc * D) * (vxd - S / Tc * xi - C * vxi)
    py = -a * (C - 1) / D * (yd - C * yi - Tc * S * vyi) - b * S / (Tc * D) * (vyd - S / Tc * yi - C * vyi)

    """Plot out all the foot placements"""
    # plt.plot(px0, py0, 'x')
    # plt.plot(px, py, 'o')


def cal_CoM_Trajectory(t):
    global sx, n
    global xi, yi, vxi, vyi, xt, yt, vxt, vyt

    """
    xt: x_tmp
    yt: y_tmp
    vxt: vx_tmp
    vyt: vy_tmp
    """
    xt = xi
    yt = yi
    vxt = vxi
    vyt = vyi

    xt = [xt, (xi - px) * np.cosh((t % Tsup) / Tc) + Tc * vxi * np.sinh((t % Tsup) / Tc) + px]
    yt = [yt, (yi - py) * np.cosh((t % Tsup) / Tc) + Tc * vyi * np.sinh((t % Tsup) / Tc) + py]

    vxt = [vxt, (xi - px) / Tc * np.sinh((t % Tsup) / Tc) + vxi * np.cosh((t % Tsup) / Tc)]
    vyt = [vyt, (yi - py) / Tc * np.sinh((t % Tsup) / Tc) + vyi * np.cosh((t % Tsup) / Tc)]

    print(px)
    print(px0)

    """Calculate the actual foot place for step n"""
    global n
    if n != int(t // Tsup):
        n = int(t // Tsup)
        if n < np.size(sx):
            cal_foot_place()
            """
            Refresh the steps.
            """
            xi = xt[-1]
            yi = yt[-1]
            vxi = vxt[-1]
            vyi = vyt[-1]

    print(xt)
    return xt[-1], vxt[-1]
